[diskann-vector] Support truly unaligned distances.#981
Open
hildebrandmw wants to merge 6 commits intomainfrom
Open
[diskann-vector] Support truly unaligned distances.#981hildebrandmw wants to merge 6 commits intomainfrom
hildebrandmw wants to merge 6 commits intomainfrom
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds first-class support in diskann-vector for computing SIMD-accelerated distances over truly under-aligned vector buffers (e.g., alignment 1), avoiding the need to copy data just to form &[T].
Changes:
- Introduces
UnalignedSlice+AsUnalignedand re-exports them from the crate root. - Updates SIMD distance kernels and specialization/dispatch plumbing to accept
AsUnalignedinputs. - Extends
Distancewithcall_unalignedand adds tests that exercise intentionally misaligned buffers.
Reviewed changes
Copilot reviewed 8 out of 9 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
diskann-vector/src/unaligned.rs |
Adds UnalignedSlice, AsUnaligned, and a test-only Buffer to create intentionally misaligned data. |
diskann-vector/src/lib.rs |
Exposes the new unaligned APIs from the crate root. |
diskann-vector/src/test_util.rs |
Refactors test harness to accept a &mut dyn DistanceChecker (trait object). |
diskann-vector/src/distance/simd.rs |
Changes simd_op to accept AsUnaligned and adds tests validating unaligned correctness/Miri safety. |
diskann-vector/src/distance/implementations.rs |
Updates architecture hooks and fixed-size specialization to operate on AsUnaligned / UnalignedSlice. |
diskann-vector/src/distance/distance_provider.rs |
Switches dispatched function signature to UnalignedSlice and adds Distance::call_unaligned. |
diskann-vector/Cargo.toml |
Adds bytemuck (dev) and enables half/bytemuck for tests. |
diskann-providers/src/model/pq/distance/multi.rs |
Adjusts reference distance calls to pass slices via explicit deref (&*...). |
Cargo.lock |
Records the new bytemuck dependency resolution. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #981 +/- ##
==========================================
- Coverage 90.63% 89.51% -1.13%
==========================================
Files 460 461 +1
Lines 85424 85549 +125
==========================================
- Hits 77427 76576 -851
- Misses 7997 8973 +976
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 8 out of 9 changed files in this pull request and generated 2 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
JordanMaples
approved these changes
Apr 29, 2026
arkrishn94
approved these changes
Apr 29, 2026
Contributor
arkrishn94
left a comment
arkrishn94
left a comment
There was a problem hiding this comment.
Looks good Mark, mostly had small comments only. The only callout is the question on the indirection through AsUnaligned. Would love to understand why this is needed.
arkrishn94
approved these changes
May 4, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
An internal user has a case where full-precision vectors (e.g.
f32) are stored in completely unaligned buffers (e.g. align of 1), requiring a data copy to align the data before the slices can be safely constructed. However, our distance function implementations useSIMDVector::load_unalignedunder the hood, which are compatible with under-aligned pointers.This PR exposes a proper API to the
DistanceProvidertrait (via theDistancetype) for invoking the SIMD implementations with unaligned pointers.Suggested Reviewing Order
diskann-wide: The implementations ofSIMDVector::load*andSIMDVector::store*already support underaligned pointers. This PR updates the documentation and restructures the load/store tests to verify this property (we were already using this property in some of the quantized distance kernels). The new load/store tests successfully pass Miri.unaligned.rs- a newUnalignedSliceis added for unaligned slices. This is just a pointer + length pair with some validity requirements but no alignment requirement. Conversions from&[T]and&[T; N]are added and the traitAsUnalignedreplaces the use ofAsRef<[T]>and the internalToSlicetraits.A test-only
Bufferis used to purposely offset simple types to exercise the unaligned cases.distance/simd.rs: Thesimd_opkernel is tweaked to acceptAsUnalignedinstead ofAsRef. Checks have been added to the existing tests to ensure that the under-unaligned versions are both Miri compatible and yield the exact same results as their properly aligned counterparts.distance/implementation.rs: The architecture hooks and specialization are changed to useAsUnaligned. I've investigated the code generation and the checks forimpl FTarget<...> for Specialize<N, F>are sufficient to trigger constant propagation and the full unrolling of small fixed-sized kernels.distance/distance_provider.rs: TheDistancetype is changed to passUnalignedSlices across the function pointer boundary rather than raw slices. We can keep the existing API for slices trivially viaAsUnaligned.Code Generation
Unfortunately, the order in which functions are code-generated seems to have changed with this PR. That said, the fixed-sized specializations I have spot-checked result in identical assembly with this PR as with main, which is to be expected.